"""
2.	分水岭算法主要用于图像的分割，如果目标物体是连接在一起的，则分割起来会很困难，此时经常采用分水岭分割算法，会得到比较好的效果。
分水岭分割算法把图像看成一幅地形图，其中，亮度比较强的区域像素值较大，亮度暗的区域像素值比较小，通过寻找汇水盆地和分水岭界线对图像进行分割。（35分）
"""
import cv2 as cv
import matplotlib.pyplot as plt
import numpy as np


def check_np_arr(arr, name):
    """Utility function for debugging."""
    print(f'{name}: type: {arr.dtype}, shape: {arr.shape}, min: {arr.min()}, max: {arr.max()}')


# ①	读取图像文件（water_coins.jpg），记为src
path = 'data/water_coins.jpg'
src = cv.imread(path, cv.IMREAD_COLOR)

# ②	备份图像文件，记为img
img = src.copy()

# ③	图像转化为灰度图
gray = cv.cvtColor(src, cv.COLOR_BGR2GRAY)

# ④	利用阈值函数对灰度图二值化
ret, bin = cv.threshold(gray, 0, 255, cv.THRESH_OTSU + cv.THRESH_BINARY_INV)

# ⑤	自定义开运算核函数kernel，并进行开运算，以消除图像中的噪点和空洞
kernel = np.ones((3, 3), dtype=np.uint8)
opening = cv.morphologyEx(bin, cv.MORPH_OPEN, kernel, iterations=2)

# ⑥	膨胀运算以标记确定的背景sure_bg
sure_bg = cv.morphologyEx(opening, cv.MORPH_CLOSE, kernel, iterations=3)

# ⑦	利用距离变换以确定前景sure_fg
dist = cv.distanceTransform(opening, cv.DIST_L2, 5)
xmax = dist.max()
ret, sure_fg = cv.threshold(dist, 0.7 * xmax, xmax, cv.THRESH_BINARY)
cv.normalize(sure_fg, sure_fg, 0, 255, cv.NORM_MINMAX)
sure_fg = np.uint8(sure_fg)

# ⑧	确定不确定区域unknown
unknown = sure_bg - sure_fg

# for marker
unknown_not = cv.bitwise_not(unknown)
ret, marker = cv.connectedComponents(unknown_not)

# ⑨	实施分水岭算法，标签图像将会被修改，边界区域的标记将变为 -1
watershed = cv.watershed(src, marker)
src[watershed == -1] = (0, 0, 255)

# ⑩	用2行4列绘出原始图、二值化灰度图、背景图、距离变换图、前景图、未知区域图、用JET颜色绘出标记图、图像分割结果图
spr = 2
spc = 4
spn = 0
plt.figure(figsize=(12, 6))
img_arr = [
    img, bin, sure_bg, dist,
    sure_fg, unknown, marker, src
]
title_arr = [
    'source', 'binary', 'background', 'dist trans',
    'foreground', 'unknown', 'JET', 'dest'
]
bgr2rgb = lambda x: cv.cvtColor(x, cv.COLOR_BGR2RGB)
arg_arr = [
    'bgr', 'gray', 'gray', 'gray',
    'gray', 'gray', 'jet', 'bgr'
]
for img_el, title, arg in zip(img_arr, title_arr, arg_arr):
    print(title)
    spn += 1
    plt.subplot(spr, spc, spn)
    plt.title(title)
    plt.axis('off')
    if 'bgr' == arg:
        img_el = bgr2rgb(img_el)
    if 'gray' == arg:
        plt.imshow(img_el, cmap='gray')
    else:
        plt.imshow(img_el)
